from asyncore import write
import os
from pickle import TRUE
import re
import shutil
import sys
from turtle import position
import cv2
root_path = r'D:\\project\\python\\datasets\\ccpd_green\\images\\changeNameTrain/'
file_name = os.listdir(root_path)
class_name = 'D:\\project\\python\\datasets\\ccpd_green\\labels\\classes.txt'
txt_dir = 'D:\\project\\python\\datasets\\ccpd_green\\labels\\train/'

Province={
   0:"皖",
   1:"沪",
   2:"津",
   3:"渝",
   4:"冀",
   5:"晋",
   6:"蒙",
   7:"辽",
   8:"吉",
   9:"黑",
   10:"苏",
   11:"浙",
   12:"京",
   13:"闽",
   14:"赣",
   15:"鲁",
   16:"豫",
   17:"鄂",
   18:"湘",
   19:"粤",
   20:"桂",
   21:"琼",
   22:"川",
   23:"贵",
   24:"云",
   25:"藏",
   26:"陕",
   27:"甘",
   28:"青",
   29:"宁",
   30:"新"
}

Letter=['A','B','C','D','E','F','G','H','J','K','L','M','N','P','Q','R','S','T','U','V','W','X','Y','Z','0','1','2','3','4','5','6','7','8','9']

def convert(size, box):
    dw = 1./size[0]
    dh = 1./size[1]
    x = (box[0] + box[1])/2.0
    y = (box[2] + box[3])/2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def writeTxt():
    if os.path.exists(class_name)==TRUE:
        os.remove(class_name)
    f = open(class_name,'w')
    f.write("placeholder\nplate\n")
    join_list = []
    write_index=2
    if not os.path.exists(txt_dir):
        os.mkdir(txt_dir)
    else:
        shutil.rmtree(txt_dir)
        os.mkdir(txt_dir)
    for i in file_name:
        print(i)
        save_name = i.split(".")[0]
        if os.path.exists(txt_dir+save_name+'.txt'):
            os.remove(txt_dir+save_name+'.txt')
        context=open(txt_dir+save_name+'.txt','w')
        img = cv2.imread(root_path+i)
        if img is None:
                print('文件%s读取失败'%i)
                continue
        height = img.shape[0]
        width = img.shape[1]
        depth = img.shape[2]
        position = i.split("-")[1]
        x1s = position.split("_")[0]
        x2s = position.split("_")[1]
        x1Left = float(x1s.split("&")[0])
        x1Up = float(x1s.split("&")[1])
        x2Left = float(x2s.split("&")[0])
        x2Up = float(x2s.split("&")[1])
        upTobottom=x2Up-x1Up
        b=(x1Left,x2Left,x1Up,x2Up)
        bb = convert((width,height),b)
        spacings = format((x2Left-x1Left)/9,'0.3f')
        context.write(str(1)+' '+str(format(bb[0],'.6f'))+' '+str(format(bb[1],'.6f'))+' '+str(format(bb[2],'.6f'))+' '+str(format(bb[3],'.6f')))
        carNumber = i.split("-")[-3]
        province = carNumber.split('_')[0]
        letters = carNumber.split('_')
        spacings = float(spacings)*1.2
        start_plate = float(x1Left)+float(spacings)
        # index=0
        provinceStr:str = Province[int(province)]
        findProvince=-1
        try:
            findProvince = join_list.index(provinceStr)
        except ValueError:
            findProvince = -1

        if findProvince==-1:
            join_list.append(provinceStr)
            #write province to txt
            write_index=write_index+1
            f.write(provinceStr+'\n')
        try:
            findProvince = join_list.index(provinceStr)
            
        except ValueError:
            findProvince = -1
        
        b=(x1Left+spacings*0.1,float(start_plate),x1Up+upTobottom*0.15,x2Up-upTobottom*0.13)
        bb = convert((width,height),b)
        context.write('\n'+str(findProvince+2)+' '+str(format(bb[0],'.6f'))+' '+str(format(bb[1],'.6f'))+' '+str(format(bb[2],'.6f'))+' '+str(format(bb[3],'.6f')))
        
        letterIndex=-1
        last_start=start_plate
        
        for _i in letters:
            letterIndex=letterIndex+1
            if letterIndex==0:
                continue
            findIndex = -1
            a=Letter[int(_i)]
            try:
                findIndex = join_list.index(a)
            except ValueError:
                findIndex = -1
            if findIndex==-1:
                join_list.append(a)
                write_index=write_index+1
                f.write(a+'\n')
           
            findIndex = join_list.index(a)
            start_plate=float(x1Left+2)+float(spacings*0.9)*(letterIndex+1)
            if letterIndex>1:
                start_plate=last_start+float(spacings*0.8)
            b=(float(last_start),float(start_plate),x1Up+upTobottom*0.15,x2Up-upTobottom*0.13)
            bb = convert((width,height),b)
            context.write('\n'+str(findIndex+2)+' '+str(format(bb[0],'.6f'))+' '+str(format(bb[1],'.6f'))+' '+str(format(bb[2],'.6f'))+' '+str(format(bb[3],'.6f')))        
            last_start=start_plate
            if letterIndex==1:
                last_start=last_start+spacings*0.7
    f.close()
    shutil.copy(class_name,txt_dir)


def readTxt():
    for i in os.listdir(txt_dir):
        f=open(txt_dir+i)
        readStr=f.read()
        if readStr.find('0 '):
            print(i,readStr)

def writeYaml(class_name:str,file_name:str):
    f = open(class_name)
    strList=f.readlines()
    yamlFile=open(file_name,'a',encoding='utf-8')
    index = 0
    for line in strList:
        str=line.replace('\n','')
        yamlFile.write("\""+str+"\",")
        index=index+1
    yamlFile.write("]")
    print(index) 

other_name = 'D:\\project\\python\\license-plate-generator-master\\labels\\classes.txt'
file_name = 'new_car.yaml'

if __name__=='__main__':
   # readTxt()
#    writeTxt()
   writeYaml(other_name,file_name)




